package utils

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

// Ws websocket接口封装
type Ws struct {
	*WsBuilder
	WsConns                  []*WsConn
	muWsConns                sync.RWMutex
	mergeSubscribeCallback   func([]interface{}) interface{}
	mergeItemMaxCount        int
	mergeUnSubscribeCallback func([]interface{}) interface{}
	mergeUnSubItemMaxCount   int

	ConnUnsubs   map[*WsConn][]interface{}
	muConnUnsubs sync.RWMutex

	curSendIdx    int
	ConnSubs      map[*WsConn][]interface{}
	muConnSubs    sync.RWMutex
	PreConnSubs   map[*WsConn][]interface{}
	muPreConnSubs sync.RWMutex
	isSubmitMode  bool
	isDialed      bool
}

// NewWs 初始化Ws
func NewWs(endpoint string, handle func(*WsConn, []byte) error) *Ws {
	ws := &Ws{
		WsBuilder:   NewWsBuilder(),
		ConnUnsubs:  make(map[*WsConn][]interface{}),
		ConnSubs:    make(map[*WsConn][]interface{}),
		PreConnSubs: make(map[*WsConn][]interface{}),
	}
	ws.WsBuilder = ws.WsBuilder.
		WsUrl(endpoint). //"wss://api.huobi.pro/ws"
		AutoReconnect().
		DecompressFunc(GzipDecompress).
		ProtoHandleFunc(handle).
		ReconnectInterval(5 * time.Second)

	ws.isDialed = false
	return ws
}

// MergeSubscribeCallback 合并订阅请求
func (ws *Ws) MergeSubscribeCallback(call func([]interface{}) interface{}, max int) {
	ws.mergeSubscribeCallback = call
	ws.mergeItemMaxCount = max
}

func (ws *Ws) MergeUnSubscribeCallback(call func([]interface{}) interface{}, max int) {
	ws.mergeUnSubscribeCallback = call
	ws.mergeUnSubItemMaxCount = max
}

// Dial 拨号
func (ws *Ws) Dial(connCount int) {
	if ws.isDialed {
		fmt.Println("already dialed!")
		return
	}
	if connCount <= 0 {
		fmt.Println("connCount must >0 ")
		return
	}
	var wait sync.WaitGroup
	var mu sync.Mutex
	wait.Add(connCount)
	for i := 0; i < connCount; i++ {
		go func() {
			conn, err := ws.WsBuilder.Build()
			if err != nil {
				fmt.Println("connect error : ", err.Error())
			}
			mu.Lock()
			ws.WsConns = append(ws.WsConns, conn)
			mu.Unlock()
			wait.Done()
		}()
	}
	wait.Wait()
	ws.curSendIdx = 0
	if ws.wsConfig.OnDialed != nil {
		ws.wsConfig.OnDialed(ws)
	}
	ws.isDialed = true
}

func (ws *Ws) IsDialed() bool {
	return ws.isDialed
}

// ResetWs 重连所有连接
func (ws *Ws) ResetWs() {
	ws.muWsConns.RLock()
	defer ws.muWsConns.RUnlock()
	for _, conn := range ws.WsConns {
		conn.ReSubscribe()
	}
}

// GetConnectCount 获取连接数
func (ws *Ws) GetConnectCount() (total, connected int) {
	ws.muWsConns.RLock()
	defer ws.muWsConns.RUnlock()
	total = len(ws.WsConns)
	connected = 0
	for _, conn := range ws.WsConns {
		if conn.isConnected {
			connected++
		}
	}
	return
}

// SetSubmitMode 设置提交模式
func (ws *Ws) SetSubmitMode(b bool) {
	ws.isSubmitMode = b
}

// Submit 提交
func (ws *Ws) Submit() error {
	if ws.isSubmitMode {
		ws.muConnSubs.Lock()
		defer ws.muConnSubs.Unlock()
		ws.muPreConnSubs.Lock()
		defer ws.muPreConnSubs.Unlock()
		ws.muWsConns.RLock()
		defer ws.muWsConns.RUnlock()
		for idx := 0; idx < len(ws.WsConns); idx++ {
			var conn = ws.WsConns[idx]
			var subs = ws.PreConnSubs[conn]
			if len(subs) == 0 {
				continue
			}
			if ws.mergeSubscribeCallback != nil {
				subList := ws.GetMergeSubList(subs)
				for _, subMerged := range subList {
					err := conn.Subscribe(subMerged)
					if err != nil {
						fmt.Println("ws Subscribe error : ", err)
					}
				}
			} else {
				for _, sub := range subs {
					err := conn.Subscribe(sub)
					if err != nil {
						fmt.Println("ws Subscribe error : ", err)
					}
				}
			}
			ws.ConnSubs[conn] = append(ws.ConnSubs[conn], subs...)
			ws.PreConnSubs[conn] = nil
		}
		return nil
	}
	return errors.New("Ws Submit error : not in submit mode")
}

func (ws *Ws) Send(msg interface{}) error {
	conn := ws.NextConn()
	if nil == conn {
		return fmt.Errorf("conns is empty!")
	}
	conn.SendJsonMessage(msg)
	return nil
}

// Subscribe 订阅
func (ws *Ws) Subscribe(sub interface{}) error {

	if !ws.isDialed {
		return errors.New("ws Subscribe error : not dialed")
	}
	conn := ws.NextConn()
	// if !ws.IsSubscribAdded(sub) {
	if !ws.isSubmitMode {
		ws.muConnSubs.Lock()
		ws.ConnSubs[conn] = append(ws.ConnSubs[conn], sub)
		ws.muConnSubs.Unlock()
		// ws.ResetSubs(conn)
		if nil != ws.mergeSubscribeCallback {
			subs := []interface{}{sub}
			subList := ws.GetMergeSubList(subs)
			for _, subMerged := range subList {
				err := conn.Subscribe(subMerged)
				if nil != err {
					return err
				}
			}
		} else {
			err := conn.Subscribe(sub)
			if nil != err {
				return err
			}
		}
	} else {
		ws.muPreConnSubs.Lock()
		ws.PreConnSubs[conn] = append(ws.PreConnSubs[conn], sub)
		ws.muPreConnSubs.Unlock()
	}
	// }

	return nil
}

func (ws *Ws) NextConn() *WsConn {
	ws.muWsConns.RLock()
	defer ws.muWsConns.RUnlock()
	idx := ws.curSendIdx
	ws.curSendIdx++
	if ws.curSendIdx >= len(ws.WsConns) {
		ws.curSendIdx = 0
	}
	return ws.WsConns[idx]
}

func (ws *Ws) AddUnsub(conn *WsConn, unsub interface{}) {
	ws.muConnUnsubs.Lock()
	defer ws.muConnUnsubs.Unlock()
	if len(ws.ConnUnsubs[conn]) > 0 {
		for _, unsubed := range ws.ConnUnsubs[conn] {
			if Compare(unsub, unsubed) {
				return
			}
		}
	}
	ws.ConnUnsubs[conn] = append(ws.ConnUnsubs[conn], unsub)
}

func (ws *Ws) RemoveUnsub(conn *WsConn, unsub interface{}) {
	ws.muConnUnsubs.Lock()
	defer ws.muConnSubs.Unlock()
	if len(ws.ConnUnsubs[conn]) > 0 {
		idx := -1
		for i, unsubed := range ws.ConnUnsubs[conn] {
			if Compare(unsub, unsubed) {
				idx = i
				break
			}
		}
		if -1 != idx {
			ws.ConnUnsubs[conn] = append(ws.ConnUnsubs[conn][:idx], ws.ConnUnsubs[conn][idx+1:]...)
		}
	}
}

func (ws *Ws) SubmitUnsubs() error {
	ws.muConnUnsubs.Lock()
	defer ws.muConnUnsubs.Unlock()
	for conn, unsubs := range ws.ConnUnsubs {
		if len(unsubs) == 0 {
			continue
		}
		if ws.mergeUnSubscribeCallback != nil {
			unsubList := ws.GetUnmergeSubList(unsubs)
			for _, unsubMerged := range unsubList {
				conn.SendJsonMessage(unsubMerged)
			}
		} else {
			for _, unsub := range unsubs {
				conn.SendJsonMessage(unsub)
			}
		}
		// fmt.Printf("unsubs : %v\n", unsubs)
	}
	ws.ConnUnsubs = make(map[*WsConn][]interface{})
	return nil
}

// Unsubscribe 取消订阅
func (ws *Ws) Unsubscribe(sub interface{}, data interface{}) error {
	bSub, idx, conn := ws.IsSubscribAdded(sub)
	if bSub {
		if !ws.isSubmitMode {
			if ws.mergeUnSubscribeCallback != nil {
				unsubs := []interface{}{data}
				unsubList := ws.GetUnmergeSubList(unsubs)
				for _, unsubMerged := range unsubList {
					conn.SendJsonMessage(unsubMerged)
				}
			} else {
				conn.SendJsonMessage(data)
			}
		} else {
			ws.AddUnsub(conn, data)
		}

		ws.muConnSubs.Lock()
		subs := ws.ConnSubs[conn]
		var changed = false
		for idx = range subs {
			if Compare(sub, subs[idx]) {
				ws.ConnSubs[conn] = append(subs[:idx], subs[idx+1:]...)
				changed = true
				break
			}
		}
		ws.muConnSubs.Unlock()
		if changed {
			ws.ResetSubs(conn)
		}
		// fmt.Printf("unsubsctibe success : %+v\n", sub)
		return nil
	}

	if ws.isSubmitMode {
		if ws.RemovePreSubs(sub) {
			return nil
		}
	}

	return fmt.Errorf("sub not find : %v", sub)
}

func (ws *Ws) ResetSubs(conn *WsConn) {

	var subsData []interface{}
	ws.muConnSubs.RLock()
	subs := ws.ConnSubs[conn]
	ws.muConnSubs.RUnlock()

	if len(subs) > 0 {
		if nil != ws.mergeSubscribeCallback {
			subsData = ws.GetMergeSubList(subs)
		} else {
			for _, sub1 := range subs {
				subsData = append(subsData, sub1)
			}
		}
	}
	conn.ResetSubs(subsData)
}

func (ws *Ws) GetUnmergeSubList(unsubs []interface{}) []interface{} {
	return ws.getMergeList(unsubs, ws.mergeUnSubscribeCallback, ws.mergeUnSubItemMaxCount)
}

func (ws *Ws) GetMergeSubList(subs []interface{}) []interface{} {
	return ws.getMergeList(subs, ws.mergeSubscribeCallback, ws.mergeItemMaxCount)
}

func (ws *Ws) getMergeList(subs []interface{}, callback func([]interface{}) interface{}, max int) []interface{} {

	length := len(subs)
	maxLength := length / max
	var ret []interface{}
	for i := 0; i < maxLength; i++ {
		ret = append(ret, callback(subs[i*max:(i+1)*max]))
	}
	if length > maxLength*max {
		ret = append(ret, callback(subs[maxLength*max:]))
	}
	// fmt.Printf("GetMergeSubList ret : %+v\n", ret)
	return ret
}

// IsSubscribAdded 查看是否已经添加订阅了
func (ws *Ws) IsSubscribAdded(sub interface{}) (bool, int, *WsConn) {
	ws.muConnSubs.RLock()
	defer ws.muConnSubs.RUnlock()
	for conn, subs := range ws.ConnSubs {
		for i, sub1 := range subs {
			if Compare(sub, sub1) {
				return true, i, conn
			}
		}
	}
	return false, -1, nil
}

func (ws *Ws) RemovePreSubs(sub interface{}) bool {
	ws.muPreConnSubs.Lock()
	defer ws.muPreConnSubs.Unlock()
	for conn, subs := range ws.PreConnSubs {
		idx := -1
		for i, sub1 := range subs {
			if Compare(sub, sub1) {
				idx = i
				break
			}
		}
		if -1 != idx {
			ws.PreConnSubs[conn] = append(subs[:idx], subs[idx+1:]...)
			return true
		}
	}
	return false
}

// Close 关闭所有连接
func (ws *Ws) Close() {
	for _, con := range ws.WsConns {
		con.CloseWs()
	}
}
